# coding=utf-8

import os
from PIL import Image
import numpy as np

file_path = "/usr/data/比赛试题/比赛数据"
trade_all_labels = []
with open(os.path.join(file_path, 'train', 'train_labels.txt')) as f:
    for i in f:
        tmp = np.zeros([10])
        tmp[int(i)] = 1
        trade_all_labels.append(tmp)


# 将图片转化为必要格式
def translate_img(im):
    im_array = np.array(im)
    return (255 - np.array([j for vec in im_array for j in vec])) / 255


# 定义一个读取trade batch的函数
def next_batch(num):
    trade_images = []
    trade_labels = []
    sample = np.random.choice(60000, num, replace=False)
    for i in sample:
        im = Image.open(os.path.join(file_path, 'train', 'TrainImage', 'TestImage_%d.bmp' % (i + 1)))
        trade_images.append(translate_img(im))
        trade_labels.append(trade_all_labels[i])
        im.close()
    return np.array(trade_images), np.array(trade_labels)


# 定义一个读取预测数据的函数
def predict_img():
    predict_images = []
    for path in os.listdir(os.path.join(file_path, 'RTestImage')):
        im = Image.open(os.path.join(file_path, 'RTestImage', path))
        predict_images.append(translate_img(im))
        im.close()
    return np.array(predict_images)


# 定义一个读取测试数据的函数
def tes_batch():
    test_images = []
    test_labels = []
    for i in range(5000):
        im = Image.open(os.path.join(file_path, 'train', 'TrainImage', 'TestImage_%d.bmp' % (i + 1 + 55000)))
        test_images.append(translate_img(im))
        test_labels.append(trade_all_labels[i + 55000])
        im.close()
    return np.array(test_images), np.array(test_labels)


for i in range(20000):
    trade_images, trade_labels = next_batch(200)
    if i % 100 == 0:
        print('step %d, trade_images.size = %d, trade_labels.size = %d' % (i, trade_images.size, trade_labels.size))
